{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameter tuning\n",
    "\n",
    "**Learning Objectives**\n",
    "1. Learn how to use `cloudml-hypertune` to report the results for Cloud hyperparameter tuning trial runs\n",
    "2. Learn how to configure the `.yaml` file for submitting  a Cloud hyperparameter tuning job\n",
    "3. Submit a hyperparameter tuning job to Vertex AI\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Let's see if we can improve upon that by tuning our hyperparameters.\n",
    "\n",
    "Hyperparameters are parameters that are set *prior* to training a model, as opposed to parameters which are learned *during* training. \n",
    "\n",
    "These include learning rate and batch size, but also model design parameters such as type of activation function and number of hidden units.\n",
    "\n",
    "Here are the four most common ways to finding the ideal hyperparameters:\n",
    "1. Manual\n",
    "2. Grid Search\n",
    "3. Random Search\n",
    "4. Bayesian Optimzation\n",
    "\n",
    "**1. Manual**\n",
    "\n",
    "Traditionally, hyperparameter tuning is a manual trial and error process. A data scientist has some intution about suitable hyperparameters which they use as a starting point, then they observe the result and use that information to try a new set of hyperparameters to try to beat the existing performance. \n",
    "\n",
    "Pros\n",
    "- Educational, builds up your intuition as a data scientist\n",
    "- Inexpensive because only one trial is conducted at a time\n",
    "\n",
    "Cons\n",
    "- Requires alot of time and patience\n",
    "\n",
    "**2. Grid Search**\n",
    "\n",
    "On the other extreme we can use grid search. Define a discrete set of values to try for each hyperparameter then try every possible combination. \n",
    "\n",
    "Pros\n",
    "- Can run hundreds of trials in parallel using the cloud\n",
    "- Gauranteed to find the best solution within the search space\n",
    "\n",
    "Cons\n",
    "- Expensive\n",
    "\n",
    "**3. Random Search**\n",
    "\n",
    "Alternatively define a range for each hyperparamter (e.g. 0-256) and sample uniformly at random from that range. \n",
    "\n",
    "Pros\n",
    "- Can run hundreds of trials in parallel using the cloud\n",
    "- Requires less trials than Grid Search to find a good solution\n",
    "\n",
    "Cons\n",
    "- Expensive (but less so than Grid Search)\n",
    "\n",
    "**4. Bayesian Optimization**\n",
    "\n",
    "Unlike Grid Search and Random Search, Bayesian Optimization takes into account information from  past trials to select parameters for future trials. The details of how this is done is beyond the scope of this notebook, but if you're interested you can read how it works here [here](https://cloud.google.com/blog/products/gcp/hyperparameter-tuning-cloud-machine-learning-engine-using-bayesian-optimization). \n",
    "\n",
    "Pros\n",
    "- Picks values intelligenty based on results from past trials\n",
    "- Less expensive because requires fewer trials to get a good result\n",
    "\n",
    "Cons\n",
    "- Requires sequential trials for best results, takes longer\n",
    "\n",
    "**Vertex AI HyperTune**\n",
    "\n",
    "Vertex AI HyperTune, powered by [Google Vizier](https://ai.google/research/pubs/pub46180), uses Bayesian Optimization by default, but [also supports](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview#search_algorithms) Grid Search and Random Search. \n",
    "\n",
    "\n",
    "When tuning just a few hyperparameters (say less than 4), Grid Search and Random Search work well, but when tunining several hyperparameters and the search space is large Bayesian Optimization is best."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT = \"<YOUR PROJECT>\"\n",
    "BUCKET = \"<YOUR BUCKET>\"\n",
    "REGION = \"<YOUR REGION>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "os.environ[\"PROJECT\"] = PROJECT\n",
    "os.environ[\"BUCKET\"] = BUCKET\n",
    "os.environ[\"REGION\"] = REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "gcloud config set project $PROJECT\n",
    "gcloud config set ai/region $REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make code compatible with Vertex AI Training Service\n",
    "In order to make our code compatible with Vertex AI Training Service we need to make the following changes:\n",
    "\n",
    "1. Upload data to Google Cloud Storage \n",
    "2. Move code into a trainer Python package\n",
    "4. Submit training job with `gcloud` to train on Vertex AI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload data to Google Cloud Storage (GCS)\n",
    "\n",
    "Cloud services don't have access to our local files, so we need to upload them to a location the Cloud servers can read from. In this case we'll use GCS.\n",
    "\n",
    "To do this run the notebook [0_export_data_from_bq_to_gcs.ipynb](./0_export_data_from_bq_to_gcs.ipynb), which will export the taxifare data from BigQuery directly into a GCS bucket. If all ran smoothly, you should be able to list the data bucket by running the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud storage ls gs://$BUCKET/taxifare/data"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Move code into python package\n",
    "\n",
    "In the [previous lab](./1_training_at_scale.ipynb), we moved our code into a python package for training on Vertex AI. Let's just check that the files are there. You should see the following files in the `taxifare/trainer` directory:\n",
    " - `__init__.py`\n",
    " - `model.py`\n",
    " - `task.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -la taxifare/trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use hyperparameter tuning in your training job you must perform the following steps:\n",
    "\n",
    " 1. Specify the hyperparameter tuning configuration for your training job by including `parameters` in the `StudySpec` of your Hyperparameter Tuning Job.\n",
    "\n",
    " 2. Include the following code in your training application:\n",
    "\n",
    "  - Parse the command-line arguments representing the hyperparameters you want to tune, and use the values to set the hyperparameters for your training trial (we already exposed these parameters as command-line arguments in the earlier lab).\n",
    "\n",
    "  - Report your hyperparameter metrics during training. Note that while you could just report the metrics at the end of training, it is better to set up a callback, to take advantage or Early Stopping.\n",
    "\n",
    "  - Read in the environment variable `$AIP_MODEL_DIR`, set by Vertex AI and containing the trial number, as our `output-dir`. As the training code will be submitted several times in a parallelized fashion, it is safer to use this variable than trying to assemble a unique id within the trainer code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modify model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./taxifare/trainer/model.py\n",
    "import datetime\n",
    "import hypertune\n",
    "import logging\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow.keras import activations\n",
    "from tensorflow.keras import callbacks\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras import models\n",
    "\n",
    "from tensorflow import feature_column as fc\n",
    "\n",
    "logging.info(tf.version.VERSION)\n",
    "\n",
    "\n",
    "CSV_COLUMNS = [\n",
    "        'fare_amount',\n",
    "        'pickup_datetime',\n",
    "        'pickup_longitude',\n",
    "        'pickup_latitude',\n",
    "        'dropoff_longitude',\n",
    "        'dropoff_latitude',\n",
    "        'passenger_count',\n",
    "        'key',\n",
    "]\n",
    "LABEL_COLUMN = 'fare_amount'\n",
    "DEFAULTS = [[0.0], ['na'], [0.0], [0.0], [0.0], [0.0], [0.0], ['na']]\n",
    "DAYS = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat']\n",
    "\n",
    "\n",
    "def features_and_labels(row_data):\n",
    "    for unwanted_col in ['key']:\n",
    "        row_data.pop(unwanted_col)\n",
    "    label = row_data.pop(LABEL_COLUMN)\n",
    "    return row_data, label\n",
    "\n",
    "\n",
    "def load_dataset(pattern, batch_size, num_repeat):\n",
    "    dataset = tf.data.experimental.make_csv_dataset(\n",
    "        file_pattern=pattern,\n",
    "        batch_size=batch_size,\n",
    "        column_names=CSV_COLUMNS,\n",
    "        column_defaults=DEFAULTS,\n",
    "        num_epochs=num_repeat,\n",
    "        shuffle_buffer_size=1000000\n",
    "    )\n",
    "    return dataset.map(features_and_labels)\n",
    "\n",
    "\n",
    "def create_train_dataset(pattern, batch_size):\n",
    "    dataset = load_dataset(pattern, batch_size, num_repeat=None)\n",
    "    return dataset.prefetch(1)\n",
    "\n",
    "\n",
    "def create_eval_dataset(pattern, batch_size):\n",
    "    dataset = load_dataset(pattern, batch_size, num_repeat=1)\n",
    "    return dataset.prefetch(1)\n",
    "\n",
    "\n",
    "def parse_datetime(s):\n",
    "    if type(s) is not str:\n",
    "        s = s.numpy().decode('utf-8')\n",
    "    return datetime.datetime.strptime(s, \"%Y-%m-%d %H:%M:%S %Z\")\n",
    "\n",
    "\n",
    "def euclidean(params):\n",
    "    lon1, lat1, lon2, lat2 = params\n",
    "    londiff = lon2 - lon1\n",
    "    latdiff = lat2 - lat1\n",
    "    return tf.sqrt(londiff*londiff + latdiff*latdiff)\n",
    "\n",
    "\n",
    "def get_dayofweek(s):\n",
    "    ts = parse_datetime(s)\n",
    "    return DAYS[ts.weekday()]\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def dayofweek(ts_in):\n",
    "    return tf.map_fn(\n",
    "        lambda s: tf.py_function(get_dayofweek, inp=[s], Tout=tf.string),\n",
    "        ts_in\n",
    "    )\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def fare_thresh(x):\n",
    "    return 60 * activations.relu(x)\n",
    "\n",
    "\n",
    "def transform(inputs, NUMERIC_COLS, STRING_COLS, nbuckets):\n",
    "    # Pass-through columns\n",
    "    transformed = inputs.copy()\n",
    "    del transformed['pickup_datetime']\n",
    "\n",
    "    feature_columns = {\n",
    "        colname: fc.numeric_column(colname)\n",
    "        for colname in NUMERIC_COLS\n",
    "    }\n",
    "\n",
    "    # Scaling longitude from range [-70, -78] to [0, 1]\n",
    "    for lon_col in ['pickup_longitude', 'dropoff_longitude']:\n",
    "        transformed[lon_col] = layers.Lambda(\n",
    "            lambda x: (x + 78)/8.0,\n",
    "            name='scale_{}'.format(lon_col)\n",
    "        )(inputs[lon_col])\n",
    "\n",
    "    # Scaling latitude from range [37, 45] to [0, 1]\n",
    "    for lat_col in ['pickup_latitude', 'dropoff_latitude']:\n",
    "        transformed[lat_col] = layers.Lambda(\n",
    "            lambda x: (x - 37)/8.0,\n",
    "            name='scale_{}'.format(lat_col)\n",
    "        )(inputs[lat_col])\n",
    "\n",
    "    # Adding Euclidean dist (no need to be accurate: NN will calibrate it)\n",
    "    transformed['euclidean'] = layers.Lambda(euclidean, name='euclidean')([\n",
    "        inputs['pickup_longitude'],\n",
    "        inputs['pickup_latitude'],\n",
    "        inputs['dropoff_longitude'],\n",
    "        inputs['dropoff_latitude']\n",
    "    ])\n",
    "    feature_columns['euclidean'] = fc.numeric_column('euclidean')\n",
    "\n",
    "    # hour of day from timestamp of form '2010-02-08 09:17:00+00:00'\n",
    "    transformed['hourofday'] = layers.Lambda(\n",
    "        lambda x: tf.strings.to_number(\n",
    "            tf.strings.substr(x, 11, 2), out_type=tf.dtypes.int32),\n",
    "        name='hourofday'\n",
    "    )(inputs['pickup_datetime'])\n",
    "    feature_columns['hourofday'] = fc.indicator_column(\n",
    "        fc.categorical_column_with_identity(\n",
    "            'hourofday', num_buckets=24))\n",
    "\n",
    "    latbuckets = np.linspace(0, 1, nbuckets).tolist()\n",
    "    lonbuckets = np.linspace(0, 1, nbuckets).tolist()\n",
    "    b_plat = fc.bucketized_column(\n",
    "        feature_columns['pickup_latitude'], latbuckets)\n",
    "    b_dlat = fc.bucketized_column(\n",
    "            feature_columns['dropoff_latitude'], latbuckets)\n",
    "    b_plon = fc.bucketized_column(\n",
    "            feature_columns['pickup_longitude'], lonbuckets)\n",
    "    b_dlon = fc.bucketized_column(\n",
    "            feature_columns['dropoff_longitude'], lonbuckets)\n",
    "    ploc = fc.crossed_column(\n",
    "            [b_plat, b_plon], nbuckets * nbuckets)\n",
    "    dloc = fc.crossed_column(\n",
    "            [b_dlat, b_dlon], nbuckets * nbuckets)\n",
    "    pd_pair = fc.crossed_column([ploc, dloc], nbuckets ** 4)\n",
    "    feature_columns['pickup_and_dropoff'] = fc.embedding_column(\n",
    "            pd_pair, 100)\n",
    "\n",
    "    return transformed, feature_columns\n",
    "\n",
    "\n",
    "def rmse(y_true, y_pred):\n",
    "    return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true)))\n",
    "\n",
    "\n",
    "def build_dnn_model(nbuckets, nnsize, lr):\n",
    "    # input layer is all float except for pickup_datetime which is a string\n",
    "    STRING_COLS = ['pickup_datetime']\n",
    "    NUMERIC_COLS = (\n",
    "            set(CSV_COLUMNS) - set([LABEL_COLUMN, 'key']) - set(STRING_COLS)\n",
    "    )\n",
    "    inputs = {\n",
    "        colname: layers.Input(name=colname, shape=(), dtype='float32')\n",
    "        for colname in NUMERIC_COLS\n",
    "    }\n",
    "    inputs.update({\n",
    "        colname: layers.Input(name=colname, shape=(), dtype='string')\n",
    "        for colname in STRING_COLS\n",
    "    })\n",
    "\n",
    "    # transforms\n",
    "    transformed, feature_columns = transform(\n",
    "        inputs, NUMERIC_COLS, STRING_COLS, nbuckets=nbuckets)\n",
    "    dnn_inputs = layers.DenseFeatures(feature_columns.values())(transformed)\n",
    "\n",
    "    x = dnn_inputs\n",
    "    for layer, nodes in enumerate(nnsize):\n",
    "        x = layers.Dense(nodes, activation='relu', name='h{}'.format(layer))(x)\n",
    "    output = layers.Dense(1, name='fare')(x)\n",
    "    \n",
    "    model = models.Model(inputs, output)\n",
    "    lr_optimizer = tf.keras.optimizers.Adam(learning_rate=lr)\n",
    "    model.compile(optimizer=lr_optimizer, loss='mse', metrics=[rmse, 'mse'])\n",
    "    \n",
    "    return model\n",
    "\n",
    "\n",
    "# TODO 1\n",
    "# Instantiate the HyperTune reporting object\n",
    "hpt = hypertune.HyperTune()\n",
    "\n",
    "# Reporting callback\n",
    "# TODO 1\n",
    "class HPTCallback(tf.keras.callbacks.Callback):\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        global hpt\n",
    "        hpt.report_hyperparameter_tuning_metric(\n",
    "            hyperparameter_metric_tag='val_rmse',\n",
    "            metric_value=logs['val_rmse'],\n",
    "            global_step=epoch)\n",
    "\n",
    "\n",
    "def train_and_evaluate(hparams):\n",
    "    batch_size = hparams['batch_size']\n",
    "    nbuckets = hparams['nbuckets']\n",
    "    lr = hparams['lr']\n",
    "    nnsize = [int(s) for s in hparams['nnsize'].split()]\n",
    "    eval_data_path = hparams['eval_data_path']\n",
    "    num_evals = hparams['num_evals']\n",
    "    num_examples_to_train_on = hparams['num_examples_to_train_on']\n",
    "    output_dir = hparams['output_dir']\n",
    "    train_data_path = hparams['train_data_path']\n",
    "\n",
    "   \n",
    "    timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S')\n",
    "    savedmodel_dir = os.path.join(output_dir, 'savedmodel')\n",
    "    model_export_path = os.path.join(savedmodel_dir, timestamp)\n",
    "    checkpoint_path = os.path.join(output_dir, 'checkpoints')\n",
    "    tensorboard_path = os.path.join(output_dir, 'tensorboard')    \n",
    "\n",
    "    if tf.io.gfile.exists(output_dir):\n",
    "        tf.io.gfile.rmtree(output_dir)\n",
    "\n",
    "    model = build_dnn_model(nbuckets, nnsize, lr)\n",
    "    logging.info(model.summary())\n",
    "\n",
    "    trainds = create_train_dataset(train_data_path, batch_size)\n",
    "    evalds = create_eval_dataset(eval_data_path, batch_size)\n",
    "\n",
    "    steps_per_epoch = num_examples_to_train_on // (batch_size * num_evals)\n",
    "\n",
    "    checkpoint_cb = callbacks.ModelCheckpoint(\n",
    "        checkpoint_path,\n",
    "        save_weights_only=True,\n",
    "        verbose=1\n",
    "    )\n",
    "    tensorboard_cb = callbacks.TensorBoard(tensorboard_path,\n",
    "                                           histogram_freq=1)\n",
    "\n",
    "    history = model.fit(\n",
    "        trainds,\n",
    "        validation_data=evalds,\n",
    "        epochs=num_evals,\n",
    "        steps_per_epoch=max(1, steps_per_epoch),\n",
    "        verbose=2,  # 0=silent, 1=progress bar, 2=one line per epoch\n",
    "        callbacks=[checkpoint_cb, tensorboard_cb, HPTCallback()]\n",
    "    )\n",
    "\n",
    "    # Exporting the model with default serving function.\n",
    "    tf.saved_model.save(model, model_export_path)    \n",
    "    return history\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modify task.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile taxifare/trainer/task.py\n",
    "import argparse\n",
    "import json\n",
    "import os\n",
    "\n",
    "from trainer import model\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\n",
    "        \"--batch_size\",\n",
    "        help = \"Batch size for training steps\",\n",
    "        type = int,\n",
    "        default = 32\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--eval_data_path\",\n",
    "        help = \"GCS location pattern of eval files\",\n",
    "        required = True\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nnsize\",\n",
    "        help = \"Hidden layer sizes (provide space-separated sizes)\",\n",
    "        default=\"32 8\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nbuckets\",\n",
    "        help = \"Number of buckets to divide lat and lon with\",\n",
    "        type = int,\n",
    "        default = 10\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--lr\",\n",
    "        help = \"learning rate for optimizer\",\n",
    "        type = float,\n",
    "        default = 0.001\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_evals\",\n",
    "        help = \"Number of times to evaluate model on eval data training.\",\n",
    "        type = int,\n",
    "        default = 5\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_examples_to_train_on\",\n",
    "        help = \"Number of examples to train on.\",\n",
    "        type = int,\n",
    "        default = 100\n",
    "    )\n",
    "    parser.add_argument(\n",
    "    \"--output_dir\",\n",
    "        help = \"GCS location to write checkpoints and export models\",\n",
    "        default = os.getenv(\"AIP_MODEL_DIR\")\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--train_data_path\",\n",
    "        help = \"GCS location pattern of train files containing eval URLs\",\n",
    "        required = True\n",
    "    )\n",
    "\n",
    "    args, _  = parser.parse_known_args()\n",
    "        \n",
    "    hparams = args.__dict__\n",
    "    print(\"output_dir\", hparams[\"output_dir\"])\n",
    "    model.train_and_evaluate(hparams)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile taxifare/setup.py\n",
    "from setuptools import find_packages\n",
    "from setuptools import setup\n",
    "\n",
    "setup(\n",
    "    name='taxifare_trainer',\n",
    "    version='0.1',\n",
    "    packages=find_packages(),\n",
    "    include_package_data=True,\n",
    "    description='Taxifare model training application.'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "cd taxifare\n",
    "python ./setup.py sdist --formats=gztar\n",
    "cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "gcloud storage cp taxifare/dist/taxifare_trainer-0.1.tar.gz gs://${BUCKET}/taxifare/"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create HyperparameterTuningJob\n",
    "\n",
    "Create a StudySpec object to hold the hyperparameter tuning configuration for your training job, and add the StudySpec to your hyperparameter tuning job.\n",
    "\n",
    "In your StudySpec `metric`, set the `metric_id` to a value representing your chosen metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Output directory and job name\n",
    "TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)\n",
    "BASE_OUTPUT_DIR=gs://${BUCKET}/taxifare_$TIMESTAMP\n",
    "JOB_NAME=taxifare_$TIMESTAMP\n",
    "echo ${BASE_OUTPUT_DIR} ${REGION} ${JOB_NAME}\n",
    "\n",
    "# Vertex AI machines to use for training\n",
    "PYTHON_PACKAGE_URI=\"gs://${BUCKET}/taxifare/taxifare_trainer-0.1.tar.gz\"\n",
    "MACHINE_TYPE=\"n1-standard-4\"\n",
    "REPLICA_COUNT=1\n",
    "PYTHON_PACKAGE_EXECUTOR_IMAGE_URI=\"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-3:latest\"\n",
    "PYTHON_MODULE=\"trainer.task\"\n",
    "\n",
    "# Model and training hyperparameters\n",
    "BATCH_SIZE=15\n",
    "NUM_EXAMPLES_TO_TRAIN_ON=100\n",
    "NUM_EVALS=10\n",
    "NBUCKETS=10\n",
    "LR=0.001\n",
    "NNSIZE=\"32 8\"\n",
    "\n",
    "# GCS paths\n",
    "GCS_PROJECT_PATH=gs://$BUCKET/taxifare\n",
    "DATA_PATH=$GCS_PROJECT_PATH/data\n",
    "TRAIN_DATA_PATH=$DATA_PATH/taxi-train*\n",
    "EVAL_DATA_PATH=$DATA_PATH/taxi-valid*\n",
    "\n",
    "\n",
    "echo > ./config.yaml \"displayName: $JOB_NAME\n",
    "studySpec:\n",
    "  metrics:\n",
    "  - metricId: val_rmse\n",
    "    goal: MINIMIZE\n",
    "  parameters:\n",
    "  - parameterId: lr\n",
    "    doubleValueSpec:\n",
    "      minValue: 0.0001\n",
    "      maxValue: 0.1\n",
    "    scaleType: UNIT_LOG_SCALE\n",
    "  - parameterId: nbuckets\n",
    "    integerValueSpec:\n",
    "      minValue: 10\n",
    "      maxValue: 25\n",
    "    scaleType: UNIT_LINEAR_SCALE\n",
    "  - parameterId: batch_size\n",
    "    discreteValueSpec:\n",
    "      values:\n",
    "      - 15\n",
    "      - 30\n",
    "      - 50\n",
    "    scaleType: UNIT_LINEAR_SCALE\n",
    "  algorithm: ALGORITHM_UNSPECIFIED # results in Bayesian optimization\n",
    "trialJobSpec:\n",
    "  baseOutputDirectory:\n",
    "    outputUriPrefix: $BASE_OUTPUT_DIR\n",
    "  workerPoolSpecs:\n",
    "  - machineSpec:\n",
    "      machineType: $MACHINE_TYPE\n",
    "    pythonPackageSpec:\n",
    "      args:\n",
    "      - --train_data_path=$TRAIN_DATA_PATH\n",
    "      - --eval_data_path=$EVAL_DATA_PATH\n",
    "      - --batch_size=$BATCH_SIZE\n",
    "      - --num_examples_to_train_on=$NUM_EXAMPLES_TO_TRAIN_ON\n",
    "      - --num_evals=$NUM_EVALS\n",
    "      - --nbuckets=$NBUCKETS\n",
    "      - --lr=$LR\n",
    "      - --nnsize=$NNSIZE\n",
    "      executorImageUri: $PYTHON_PACKAGE_EXECUTOR_IMAGE_URI\n",
    "      packageUris:\n",
    "      - $PYTHON_PACKAGE_URI\n",
    "      pythonModule: $PYTHON_MODULE\n",
    "    replicaCount: $REPLICA_COUNT\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)\n",
    "JOB_NAME=taxifare_$TIMESTAMP\n",
    "echo $REGION\n",
    "echo $JOB_NAME\n",
    "gcloud beta ai hp-tuning-jobs create \\\n",
    "    --region=$REGION \\\n",
    "    --display-name=$JOB_NAME \\\n",
    "    --config=config.yaml \\\n",
    "    --max-trial-count=10 \\\n",
    "    --parallel-trial-count=2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could have also used the Vertex AI Python SDK to achieve the same, as below:"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from datetime import datetime\n",
    "from google.cloud import aiplatform\n",
    "\n",
    "\n",
    "# Output directory and jobID\n",
    "timestamp_str=datetime.strftime(datetime.now(), '%Y%m%d_%H%M%S')\n",
    "BASE_OUTPUT_DIR=f\"gs://{BUCKET}/taxifare_{timestamp_str}\"\n",
    "JOB_NAME=f\"taxifare_{timestamp_str}\"\n",
    "print(BASE_OUTPUT_DIR, REGION, JOB_NAME)\n",
    "\n",
    "\n",
    "# Vertex AI machines to use for training\n",
    "PYTHON_PACKAGE_URIS=f\"gs://{BUCKET}/taxifare/taxifare_trainer-0.1.tar.gz\"\n",
    "MACHINE_TYPE=\"n1-standard-4\"\n",
    "REPLICA_COUNT=1\n",
    "PYTHON_PACKAGE_EXECUTOR_IMAGE_URI=\"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-3:latest\"\n",
    "PYTHON_MODULE=\"trainer.task\"\n",
    "\n",
    "# Model and training hyperparameters\n",
    "BATCH_SIZE=15\n",
    "NUM_EXAMPLES_TO_TRAIN_ON=100\n",
    "NUM_EVALS=10\n",
    "NBUCKETS=10\n",
    "LR=0.001\n",
    "NNSIZE=\"32 8\"\n",
    "\n",
    "\n",
    "# GCS paths.\n",
    "GCS_PROJECT_PATH=f\"gs://{BUCKET}/taxifare\"\n",
    "DATA_PATH=f\"{GCS_PROJECT_PATH}/data\"\n",
    "TRAIN_DATA_PATH=f\"{DATA_PATH}/taxi-train*\"\n",
    "EVAL_DATA_PATH=f\"{DATA_PATH}/taxi-valid*\"\n",
    "\n",
    "# custom container\n",
    "IMAGE_NAME=\"taxifare_training_container\"\n",
    "IMAGE_URI=f\"gcr.io/{PROJECT}/{IMAGE_NAME}\"\n",
    "\n",
    "def create_hyperparameter_tuning_job_python_package_sample(\n",
    "    project: str,\n",
    "    display_name: str,\n",
    "    executor_image_uri: str,\n",
    "    package_uri: str,\n",
    "    python_module: str,\n",
    "    location: str = REGION,\n",
    "    api_endpoint: str = f\"{REGION}-aiplatform.googleapis.com\",\n",
    "):\n",
    "    # The Vertex AI services require regional API endpoints.\n",
    "    client_options = {\"api_endpoint\": api_endpoint}\n",
    "    # Initialize client that will be used to create and send requests.\n",
    "    # This client only needs to be created once, and can be reused for multiple requests.\n",
    "    client = aiplatform.gapic.JobServiceClient(client_options=client_options)\n",
    "\n",
    "    \n",
    "    # study_spec\n",
    "    metric = {\n",
    "        \"metric_id\": \"val_rmse\",\n",
    "        \"goal\": aiplatform.gapic.StudySpec.MetricSpec.GoalType.MINIMIZE,\n",
    "    }\n",
    "    parameter_lr = {\n",
    "        \"parameter_id\": \"lr\",\n",
    "        \"double_value_spec\": {\"min_value\": 0.0001, \"max_value\": 0.1},\n",
    "        \"scale_type\": aiplatform.gapic.StudySpec.ParameterSpec.ScaleType.UNIT_LOG_SCALE,\n",
    "    }    \n",
    "    parameter_nbuckets = {\n",
    "        \"parameter_id\": \"nbuckets\",\n",
    "        \"integer_value_spec\": {\"min_value\": 10, \"max_value\": 25},\n",
    "        \"scale_type\": aiplatform.gapic.StudySpec.ParameterSpec.ScaleType.UNIT_LINEAR_SCALE,\n",
    "    }    \n",
    "    parameter_batchsize = {\n",
    "        \"parameter_id\": \"batch_size\",\n",
    "        \"discrete_value_spec\": {\"values\": [15, 30, 50]},\n",
    "        \"scale_type\": aiplatform.gapic.StudySpec.ParameterSpec.ScaleType.UNIT_LINEAR_SCALE,\n",
    "    }    \n",
    "\n",
    "    # trial_job_spec\n",
    "    worker_pool_spec = {\n",
    "        \"machine_spec\": {\n",
    "            \"machine_type\": \"n1-standard-4\",\n",
    "        },\n",
    "        \"replica_count\": 1,\n",
    "        \"python_package_spec\": {\n",
    "            \"executor_image_uri\": executor_image_uri,\n",
    "            \"package_uris\": [package_uri],\n",
    "            \"python_module\": python_module,\n",
    "            \"args\": [\n",
    "                            f\"--eval_data_path={EVAL_DATA_PATH}\",\n",
    "                            f\"--train_data_path={TRAIN_DATA_PATH}\",\n",
    "                            f\"--batch_size={BATCH_SIZE}\",\n",
    "                            f\"--num_examples_to_train_on={NUM_EXAMPLES_TO_TRAIN_ON}\",\n",
    "                            f\"--num_evals={NUM_EVALS}\",\n",
    "                            f\"--nbuckets={NBUCKETS}\",\n",
    "                            f\"--lr={LR}\",\n",
    "                            f\"--nnsize={NNSIZE}\"\n",
    "            ],\n",
    "        },\n",
    "    }\n",
    "\n",
    "    # hyperparameter_tuning_job\n",
    "    hyperparameter_tuning_job = {\n",
    "        \"display_name\": display_name,\n",
    "        \"max_trial_count\": 10,\n",
    "        \"parallel_trial_count\": 2,\n",
    "        \"study_spec\": {\n",
    "            \"metrics\": [metric],\n",
    "            \"parameters\": [\n",
    "                parameter_lr,\n",
    "                parameter_nbuckets,\n",
    "                parameter_batchsize,\n",
    "            ],\n",
    "            \"algorithm\": aiplatform.gapic.StudySpec.Algorithm.ALGORITHM_UNSPECIFIED,  # results in Bayesian optimization\n",
    "#            \"median_automated_stopping_spec\": {} # early stopping: only available in v1beta1 as of writing\n",
    "        },\n",
    "        \"trial_job_spec\": {\n",
    "            \"worker_pool_specs\": [worker_pool_spec],\n",
    "            \"base_output_directory\": {\n",
    "                'output_uri_prefix': BASE_OUTPUT_DIR,\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "    parent = f\"projects/{project}/locations/{location}\"\n",
    "    response = client.create_hyperparameter_tuning_job(parent=parent, hyperparameter_tuning_job=hyperparameter_tuning_job)\n",
    "    print(\"response:\", response)\n",
    "\n",
    "create_hyperparameter_tuning_job_python_package_sample(\n",
    "    project=PROJECT,\n",
    "    display_name=JOB_NAME,\n",
    "    executor_image_uri=PYTHON_PACKAGE_EXECUTOR_IMAGE_URI,\n",
    "    package_uri=PYTHON_PACKAGE_URIS,\n",
    "    python_module=PYTHON_MODULE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021 Google Inc. Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-gpu.2-3.m71",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-3:m71"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
